"""MCP服务器主类"""

import asyncio
from typing import Dict, Any, Optional
from loguru import logger

from .protocol import MCPProtocol
import sys
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from src.core.context_orchestrator import ContextOrchestrator


class MCPServer:
    """MCP服务器主类"""
    
    def __init__(self, orchestrator: Optional[ContextOrchestrator], config: Dict[str, Any]):
        self.orchestrator = orchestrator
        self.config = config
        self.mcp_config = config.get('mcp', {})
        
        # 初始化协议处理器
        self.protocol = MCPProtocol(orchestrator, config)
        
        # 服务器状态
        self.is_running = False
        self.connections = {}
        
        logger.info("MCP服务器已初始化")
    
    async def start(self):
        """启动MCP服务器"""
        if self.is_running:
            logger.warning("MCP服务器已在运行")
            return
        
        try:
            # 如果有orchestrator，确保数据库连接已初始化
            if self.orchestrator and not self.orchestrator.db_manager.is_initialized:
                await self.orchestrator.db_manager.initialize()
            
            self.is_running = True
            logger.info("MCP服务器已启动")
            
        except Exception as e:
            logger.error(f"启动MCP服务器失败: {e}")
            raise
    
    async def stop(self):
        """停止MCP服务器"""
        if not self.is_running:
            logger.warning("MCP服务器未在运行")
            return
        
        try:
            # 关闭所有连接
            for connection_id in list(self.connections.keys()):
                await self.disconnect_client(connection_id)
            
            self.is_running = False
            logger.info("MCP服务器已停止")
            
        except Exception as e:
            logger.error(f"停止MCP服务器失败: {e}")
    
    async def handle_message(self, message: str, connection_id: Optional[str] = None) -> str:
        """处理客户端消息"""
        if not self.is_running:
            raise RuntimeError("MCP服务器未运行")
        
        try:
            # 如果没有连接ID，生成一个
            if connection_id is None:
                import uuid
                connection_id = str(uuid.uuid4())
            
            # 记录连接
            if connection_id not in self.connections:
                await self.connect_client(connection_id)
            
            # 更新连接活动时间
            self.connections[connection_id]['last_activity'] = asyncio.get_event_loop().time()
            
            # 处理消息
            response = await self.protocol.handle_message(message, connection_id)
            
            logger.debug(f"处理消息完成: {connection_id}")
            return response
            
        except Exception as e:
            logger.error(f"处理消息失败: {e}")
            # 返回错误响应
            import json
            error_response = {
                'jsonrpc': '2.0',
                'error': {
                    'code': -32603,
                    'message': 'Internal error',
                    'data': str(e)
                },
                'id': None
            }
            return json.dumps(error_response, ensure_ascii=False)
    
    async def connect_client(self, connection_id: str, client_info: Optional[Dict[str, Any]] = None):
        """连接客户端"""
        try:
            self.connections[connection_id] = {
                'connected_at': asyncio.get_event_loop().time(),
                'last_activity': asyncio.get_event_loop().time(),
                'client_info': client_info or {},
                'message_count': 0
            }
            
            logger.info(f"客户端已连接: {connection_id}")
            
        except Exception as e:
            logger.error(f"连接客户端失败: {e}")
            raise
    
    async def disconnect_client(self, connection_id: str):
        """断开客户端连接"""
        try:
            if connection_id in self.connections:
                connection_info = self.connections[connection_id]
                duration = asyncio.get_event_loop().time() - connection_info['connected_at']
                message_count = connection_info['message_count']
                
                del self.connections[connection_id]
                
                logger.info(f"客户端已断开: {connection_id} (持续时间: {duration:.2f}s, 消息数: {message_count})")
            
        except Exception as e:
            logger.error(f"断开客户端连接失败: {e}")
    
    async def broadcast_notification(self, method: str, params: Dict[str, Any]):
        """向所有连接的客户端广播通知"""
        import json
        
        notification = {
            'jsonrpc': '2.0',
            'method': method,
            'params': params
        }
        
        notification_str = json.dumps(notification, ensure_ascii=False)
        
        # 这里需要具体的传输层实现来发送通知
        # 在Stdio和WebSocket处理器中会实现具体的发送逻辑
        logger.debug(f"广播通知: {method} 到 {len(self.connections)} 个连接")
    
    def get_connection_info(self, connection_id: str) -> Optional[Dict[str, Any]]:
        """获取连接信息"""
        return self.connections.get(connection_id)
    
    def get_all_connections(self) -> Dict[str, Dict[str, Any]]:
        """获取所有连接信息"""
        return self.connections.copy()
    
    async def cleanup_inactive_connections(self, max_inactive_seconds: int = 3600):
        """清理不活跃的连接"""
        current_time = asyncio.get_event_loop().time()
        inactive_connections = []
        
        for connection_id, connection_info in self.connections.items():
            if current_time - connection_info['last_activity'] > max_inactive_seconds:
                inactive_connections.append(connection_id)
        
        for connection_id in inactive_connections:
            await self.disconnect_client(connection_id)
        
        # 同时清理协议层的不活跃会话
        cleaned_sessions = self.protocol.cleanup_inactive_sessions()
        
        if inactive_connections or cleaned_sessions:
            logger.info(f"清理了 {len(inactive_connections)} 个不活跃连接和 {cleaned_sessions} 个会话")
        
        return len(inactive_connections)
    
    async def get_server_stats(self) -> Dict[str, Any]:
        """获取服务器统计信息"""
        try:
            current_time = asyncio.get_event_loop().time()
            
            # 连接统计
            total_connections = len(self.connections)
            active_connections = 0
            total_messages = 0
            
            for connection_info in self.connections.values():
                if current_time - connection_info['last_activity'] < 300:  # 5分钟内活跃
                    active_connections += 1
                total_messages += connection_info['message_count']
            
            # 数据库健康状态
            db_health = await self.orchestrator.db_manager.health_check()
            
            # 协议统计
            protocol_sessions = len(self.protocol.sessions)
            
            return {
                'server_status': 'running' if self.is_running else 'stopped',
                'connections': {
                    'total': total_connections,
                    'active': active_connections,
                    'total_messages': total_messages
                },
                'protocol': {
                    'sessions': protocol_sessions,
                    'tools_registered': len(self.protocol.tools),
                    'resources_registered': len(self.protocol.resources)
                },
                'database_health': db_health,
                'uptime_seconds': current_time - getattr(self, '_start_time', current_time)
            }
            
        except Exception as e:
            logger.error(f"获取服务器统计失败: {e}")
            return {'error': str(e)}
    
    async def health_check(self) -> Dict[str, Any]:
        """健康检查"""
        try:
            # 检查服务器状态
            server_healthy = self.is_running
            
            # 检查数据库连接
            db_health = await self.orchestrator.db_manager.health_check()
            db_healthy = all(db_health.values())
            
            # 检查内存使用
            import psutil
            process = psutil.Process()
            memory_info = process.memory_info()
            memory_usage_mb = memory_info.rss / 1024 / 1024
            
            # 检查连接数
            connection_count = len(self.connections)
            max_connections = self.mcp_config.get('max_connections', 1000)
            connections_healthy = connection_count < max_connections * 0.9
            
            overall_healthy = server_healthy and db_healthy and connections_healthy
            
            return {
                'healthy': overall_healthy,
                'status': 'healthy' if overall_healthy else 'unhealthy',
                'checks': {
                    'server': server_healthy,
                    'database': db_healthy,
                    'connections': connections_healthy
                },
                'metrics': {
                    'memory_usage_mb': memory_usage_mb,
                    'connection_count': connection_count,
                    'max_connections': max_connections
                },
                'database_status': db_health,
                'timestamp': asyncio.get_event_loop().time()
            }
            
        except Exception as e:
            logger.error(f"健康检查失败: {e}")
            return {
                'healthy': False,
                'status': 'error',
                'error': str(e),
                'timestamp': asyncio.get_event_loop().time()
            }
    
    def register_tool(self, name: str, func, schema: Optional[Dict[str, Any]] = None):
        """注册自定义工具"""
        self.protocol.register_tool(name, func, schema)
    
    def register_resource(self, name: str, func):
        """注册自定义资源"""
        self.protocol.register_resource(name, func)
    
    async def __aenter__(self):
        """异步上下文管理器入口"""
        await self.start()
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器出口"""
        await self.stop()